{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction to polarization\n",
    "\n",
    "We will introduce how polarization is simulated in HCIPy, both for full and partially polarized light. We will show which polarization optical elements are implemented and how these are used.\n",
    "\n",
    "We'll start by implementing the required packages and set up a circular pupil."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from hcipy import *\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "pupil_grid = make_pupil_grid(256)\n",
    "\n",
    "aperture = make_magellan_aperture(True)(pupil_grid)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Types of polarized Wavefronts\n",
    "\n",
    "HCIPy supports three different types of wavefronts, which indicate different types of polarization. These are:\n",
    "\n",
    "1. Scalar wavefronts\n",
    "2. Fully polarized wavefronts\n",
    "3. Partially polarized wavefronts\n",
    "\n",
    "Each of these different types of wavefronts are implemented in a `Wavefront` object. There are a number of ways to get the polarization state of the `Wavefront`. These include (components of) the Stokes vector, the degree of linear/circular polarization, etc... The different types of wavefronts are distinguished by the tensor shape of the `electric_field` that is in the `Wavefront`.\n",
    "\n",
    "When you are not using polarization, you are most likely mostly using scalar wavefronts. These contain a scalar `electric_field`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scalar_wavefront = Wavefront(aperture)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Fully polarized wavefronts use Jones vectors as their electric fields. This means that the electric field will be a vector `Field` with two components:\n",
    "$$\\begin{bmatrix}E_x \\\\ E_y\\end{bmatrix}$$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "jones_vector = Field([aperture, 1j * aperture], pupil_grid)\n",
    "\n",
    "fully_polarized_wavefront = Wavefront(jones_vector)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Partially polarized wavefronts use Jones matrices as their electric fields. The electric field will be a 2x2 tensor `Field`:$$\\begin{bmatrix} J_{xx} & J_{xy} \\\\ J_{yx} & J_{yy}\\end{bmatrix}.$$\n",
    "\n",
    "These Jones matrices describe the change in polarization state and also phase and amplitude, from the initial wavefront to the current wavefront. The polarization state itself is set as the Stokes vector of the initial wavefront. Note that, in this way, HCIPy doesn't support spatially varying degrees of polarization.\n",
    "\n",
    "While you can create a full Jones matrix `Field` and pass that as the electric field argument of the `Wavefront`, like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "jones_matrix = Field([[aperture, aperture * 1j], [aperture * 1j, aperture]], pupil_grid)\n",
    "stokes_vector = [1, 0.3, -0.2, 0.1]\n",
    "\n",
    "partially_polarized_wavefront = Wavefront(jones_matrix, input_stokes_vector=stokes_vector)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "in most cases, we are trying to create a `Wavefront` at the start of our optical system. In this case we can pass a scalar electric field, but add an input Stokes vector as well. This creates a Jones matrix `Field` which has for each pixel an identity matrix multiplied by the scalar wavefront at that position. For example, creating a partially polarized `Wavefront` at our telescope pupil, we would do:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stokes_vector = [1, 0.3, -0.2, 0.1]\n",
    "\n",
    "partially_polarized_wavefront = Wavefront(aperture, input_stokes_vector=stokes_vector)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Polarization attributes of Wavefronts\n",
    "\n",
    "We can look at the Stokes components of the wavefront, and show them:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def imshow_stokes_vector(wavefront):\n",
    "    I_max = wavefront.I.max()\n",
    "    \n",
    "    plt.subplot(2, 2, 1)\n",
    "    plt.title('I')\n",
    "    imshow_field(wavefront.I, cmap='inferno')\n",
    "    plt.colorbar()\n",
    "    \n",
    "    plt.subplot(2, 2, 2)\n",
    "    plt.title('Q')\n",
    "    imshow_field(wavefront.Q, cmap='RdBu', vmin=-I_max, vmax=I_max)\n",
    "    plt.colorbar()\n",
    "    \n",
    "    plt.subplot(2, 2, 3)\n",
    "    plt.title('U')\n",
    "    imshow_field(wavefront.U, cmap='RdBu', vmin=-I_max, vmax=I_max)\n",
    "    plt.colorbar()\n",
    "    \n",
    "    plt.subplot(2, 2, 4)\n",
    "    plt.title('V')\n",
    "    imshow_field(wavefront.V, cmap='RdBu', vmin=-I_max, vmax=I_max)\n",
    "    plt.colorbar()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.suptitle('Scalar wavefront')\n",
    "imshow_stokes_vector(scalar_wavefront)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.suptitle('Fully polarized wavefront')\n",
    "imshow_stokes_vector(fully_polarized_wavefront)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.suptitle('Partially polarized wavefront')\n",
    "imshow_stokes_vector(partially_polarized_wavefront)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also show the polarization fraction:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(nrows=1, ncols=3)\n",
    "\n",
    "axes[0].set_title('Scalar wavefront')\n",
    "imshow_field(scalar_wavefront.degree_of_polarization, vmin=0, vmax=1, cmap='inferno', mask=aperture, ax=axes[0])\n",
    "\n",
    "axes[1].set_title('Fully polarized wavefront')\n",
    "imshow_field(fully_polarized_wavefront.degree_of_polarization, vmin=0, vmax=1, cmap='inferno', mask=aperture, ax=axes[1])\n",
    "\n",
    "axes[2].set_title('Partially polarized wavefront')\n",
    "imshow_field(partially_polarized_wavefront.degree_of_polarization, vmin=0, vmax=1, cmap='inferno', mask=aperture, ax=axes[2])\n",
    "\n",
    "plt.colorbar(ax=axes.tolist(), orientation=\"horizontal\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see that the scalar wavefront is considered an unpolarized wavefront, the fully polarized wavefront is 100% polarized, and the partially polarized wavefront permits a partial polarization fraction, which is exactly as we expected."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Optical propagation with polarized wavefronts\n",
    "\n",
    "Optical propagations through optical elements with polarized wavefronts are seamless. Take for example a `FraunhoferPropagator` propagation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "focal_grid = make_focal_grid(3, 8)\n",
    "prop = FraunhoferPropagator(pupil_grid, focal_grid)\n",
    "\n",
    "scalar_image = prop(scalar_wavefront)\n",
    "fully_polarized_image = prop(fully_polarized_wavefront)\n",
    "partially_polarized_image = prop(partially_polarized_wavefront)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.suptitle('Partially polarized wavefront')\n",
    "imshow_stokes_vector(partially_polarized_image)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As the `FraunhoferPropagator` does not act differently depending on the polarization of the incoming light, we see a normal Airy core which has the same polarization state as we put in. Note that the call to the propagator does not change depending on the type of the `Wavefront` that we propagate through it. All calls look exactly the same. This means that if we create a full optical system, consisting of many optical elements in a row, that any `Wavefront` can be propgated through, and we do not need to worry about having to rewrite the optical system when we want to switch to (partially) polarized light.\n",
    "\n",
    "HCIPy contains many polarization optical elements, such as half-wave plates, quarter-wave plates, general phase retarders, polarizers, polarizing beam splitters and more. Let's simulate a high-contrast imaging scenario of an unpolarized star and a partially polarized planet next to it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stokes_vector_planet = [1, 0.1, 0, 0]\n",
    "contrast = 3e-4\n",
    "angular_separation = 5 # lambda / D\n",
    "num_photons_star = 1e9\n",
    "\n",
    "wf_star = Wavefront(aperture)\n",
    "wf_star.total_power = num_photons_star\n",
    "\n",
    "e_planet = aperture * np.exp(2j * np.pi * pupil_grid.x * angular_separation)\n",
    "wf_planet = Wavefront(e_planet, input_stokes_vector=stokes_vector_planet)\n",
    "wf_planet.total_power = num_photons_star * contrast\n",
    "\n",
    "# Calculate science image for the star and planet\n",
    "science_image = prop(wf_star).power\n",
    "science_image += prop(wf_planet).power\n",
    "\n",
    "# Simulate photon noise\n",
    "science_image = large_poisson(science_image)\n",
    "\n",
    "imshow_field(np.log10(science_image / science_image.max() + 1e-20), vmin=-5, vmax=0)\n",
    "plt.colorbar()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this intensity image it is impossible to see the planet. It is completely buried below the light from the star. We need to look at a polarized light image. The simplest way, in fact to the point of oversimplification, is to split the light with a linear polarizing beam splitter and look at the difference between the two images. If the planet is polarized, it will show up in the difference image.\n",
    "\n",
    "Let's create a linear polarizing beam splitter object, and propagate the star and planet `Wavefront`s through it. This creates two wavefronts, one for each port of the beam splitter. The image for each port is calculated."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pbs = LinearPolarizingBeamSplitter(0)\n",
    "\n",
    "img_star_1, img_star_2 = pbs(prop(wf_star))\n",
    "img_planet_1, img_planet_2 = pbs(prop(wf_planet))\n",
    "\n",
    "science_image_1 = img_star_1.power + img_planet_1.power\n",
    "science_image_2 = img_star_2.power + img_planet_2.power\n",
    "\n",
    "science_image_1 = large_poisson(science_image_1)\n",
    "science_image_2 = large_poisson(science_image_2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now show these two images and their difference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(figsize=(10, 6), nrows=1, ncols=2)\n",
    "\n",
    "axes[0].set_title('Image for port 1')\n",
    "imshow_field(np.log10(science_image_1 / science_image_1.max() + 1e-20), vmin=-5, vmax=0, ax=axes[0])\n",
    "\n",
    "axes[1].set_title('Image for port 2')\n",
    "imshow_field(np.log10(science_image_2 / science_image_2.max() + 1e-20), vmin=-5, vmax=0, ax=axes[1])\n",
    "\n",
    "plt.colorbar(ax=axes.tolist(), orientation=\"horizontal\")\n",
    "plt.show()\n",
    "\n",
    "plt.title('Difference')\n",
    "imshow_field((science_image_1 - science_image_2) / science_image_1.max(), cmap='RdBu', vmin=-1e-4, vmax=1e-4)\n",
    "plt.annotate('planet', xy=(5, -1), xytext=(3, -7), arrowprops={'arrowstyle': '-|>'})\n",
    "plt.colorbar()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the difference image we can clearly see the planet light up. It has a contrast of $3\\times10^{-5}$, which is partly due to the contrast of the planet ($3\\times10^{-4}$) and the degree of polarization of the planet (10%). While this is clearly a simplified example (it can only detect the Q-component of the planet), it shows the simplicity and power of HCIPy for propagating polarized light.\n",
    "\n",
    "This introduction falls short of the many other optical elements that act on polarization. Other tutorials showcase some more of these optical elements. These include, but are not limited to, [a simple polarimeter](../SimplePolarimeter/SimplePolarimeter.rst) and [imaging with a vector apodizing phase plate coronagraph](../VectorApodizingPhasePlate/VectorApodizingPhasePlate.rst)."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  },
  "level": "beginner",
  "thumbnail_figure_index": -1
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
